#include "pch.h"
#include "AudioSampleProcessor.h"

using namespace Microsoft::Media::FragmentedMediaSink;

AudioSampleProcessor::AudioSampleProcessor(IVector<IFragmentController^>^ fragmentControllers, IMFMediaType* spInputMediaType, unsigned int inputStreamIndex)
	:ISampleProcessor(fragmentControllers, spInputMediaType, inputStreamIndex)
	, _minDuration(9223372036854775807)
	, _currentFragmentDuration(0)
	, _currentFragmentOffset(0)
	, _fragmentEndTime(0)
	, _nextSampleTime(0)
	, _isInitialized(false)
	, _samplesInFragment(0)
	, _fragmentIndex(0)
{
	for each (auto fragmentController in fragmentControllers)
	{
		_maxFragmentDuration = fragmentController->FragmentSize.Duration;
		break;
	}

	HRESULT hr = S_OK;
	if (SUCCEEDED(hr))
		hr = spInputMediaType->GetUINT32(MF_MT_AUDIO_NUM_CHANNELS, &_channels);
	if (SUCCEEDED(hr))
		hr = spInputMediaType->GetUINT32(MF_MT_AUDIO_SAMPLES_PER_SECOND, &_sampleRate);
	if (SUCCEEDED(hr))
		hr = spInputMediaType->GetUINT32(MF_MT_AUDIO_BITS_PER_SAMPLE, &_bitsPerSample);
}

AudioSampleProcessor::~AudioSampleProcessor()
{
}

HRESULT AudioSampleProcessor::ProcessSample(_In_ IMFSample* spSample)
{
	HRESULT hr = S_OK;

	if (!_isInitialized)
	{
		hr = InitializeFragments();
	}

	long long sampleTime;
	if (SUCCEEDED(hr))
	{
		hr = spSample->GetSampleTime(&sampleTime);
	}

	long long sampleDuration;
	if (SUCCEEDED(hr))
	{
		hr = spSample->GetSampleDuration(&sampleDuration);
	}

	if (SUCCEEDED(hr))
	{
		if (sampleDuration < _minDuration) _minDuration = sampleDuration;
		long long halfDuration = sampleDuration / 2;
		while ((sampleTime > _nextSampleTime + halfDuration) && (_nextSampleTime + _minDuration < _fragmentEndTime))
		{
			// too late, add a blank sample
			ComPtr<IMFSample> blankSample;
			hr = CreateBlankSample(&blankSample, _nextSampleTime, _minDuration);
			if (SUCCEEDED(hr))
			{
				hr = WriteSample(blankSample.Get());
			}
			if (FAILED(hr)) break;
		}

		if (SUCCEEDED(hr))
		{
			if (sampleTime >= _nextSampleTime - halfDuration || _samplesInFragment == 0) // else too soon, drop sample
			{
				hr = spSample->SetSampleTime(_nextSampleTime);

				if (SUCCEEDED(hr))
				{
					long long endTime = _nextSampleTime + sampleDuration;
					if (endTime > _fragmentEndTime)
					{
						// fill end of fragment with blank sample
						ComPtr<IMFSample> blankSample;
						hr = CreateBlankSample(&blankSample, _nextSampleTime, _fragmentEndTime - _nextSampleTime);
						if (SUCCEEDED(hr))
						{
							hr = WriteSample(blankSample.Get());
						}
						if (SUCCEEDED(hr))
						{
							hr = TerminateFragments();
						}
						if (SUCCEEDED(hr))
						{
							hr = ProcessSample(spSample);
						}
					}
					else
					{
						hr = WriteSample(spSample);
						if (SUCCEEDED(hr))
						{
							if (endTime == _fragmentEndTime)
							{
								hr = TerminateFragments();
							}
						}
					}
				}
			}
		}
	}

	return hr;
}

HRESULT AudioSampleProcessor::CreateBlankBuffer(_Out_ IMFMediaBuffer** ppBuffer, long long duration)
{
	HRESULT hr = S_OK;

	ComPtr<IMFMediaBuffer> spBuffer;
	*ppBuffer = nullptr;

	unsigned int sampleCount = (unsigned int)((long long)_sampleRate * duration / 10000000);
	DWORD cbLength = _channels * (_bitsPerSample / 8) * sampleCount;

	if (SUCCEEDED(hr))
		hr = MFCreateMemoryBuffer(cbLength, &spBuffer);

	if (SUCCEEDED(hr))
	{
		unsigned char *pbBuffer = nullptr;
		hr = spBuffer->Lock(&pbBuffer, NULL, NULL);

		memset(pbBuffer, 0, (size_t)cbLength);

		if (SUCCEEDED(hr))
			hr = spBuffer->SetCurrentLength(cbLength);

		if (SUCCEEDED(hr))
		{
			hr = spBuffer->Unlock();
		}
	}

	*ppBuffer = spBuffer.Detach();

	return hr;
}

HRESULT AudioSampleProcessor::CreateBlankSample(_Out_ IMFSample** ppSample, _In_ long long sampleTime, _In_ long long sampleDuration)
{
	HRESULT hr = S_OK;

	ComPtr<IMFSample> spSample;
	*ppSample = nullptr;
	hr = MFCreateSample(&spSample);

	if (SUCCEEDED(hr))
	{
		ComPtr<IMFMediaBuffer> spBuffer;

		hr = CreateBlankBuffer(&spBuffer, sampleDuration);

		if (SUCCEEDED(hr))
			hr = spSample->AddBuffer(spBuffer.Get());
	}

	if (SUCCEEDED(hr))
	{
		hr = spSample->SetSampleTime(sampleTime);
	}

	if (SUCCEEDED(hr))
	{
		hr = spSample->SetSampleDuration(sampleDuration);
	}

	*ppSample = spSample.Detach();
	return hr;
}

HRESULT AudioSampleProcessor::SendTick(_In_ long long timeStamp)
{
	return ISampleProcessor::SendTick(timeStamp - _currentFragmentOffset);
}

HRESULT AudioSampleProcessor::WriteSample(_In_ IMFSample* spSample)
{
	HRESULT hr = S_OK;

	long long sampleTime;
	if (SUCCEEDED(hr))
	{
		hr = spSample->GetSampleTime(&sampleTime);
	}

	long long sampleDuration;
	if (SUCCEEDED(hr))
	{
		hr = spSample->GetSampleDuration(&sampleDuration);
	}

	if (SUCCEEDED(hr))
	{
		_nextSampleTime = sampleTime + sampleDuration;
		_currentFragmentDuration = _nextSampleTime - _currentFragmentOffset;

		hr = spSample->SetSampleTime(_currentFragmentDuration - sampleDuration);

		if (SUCCEEDED(hr))
		{
			ISampleProcessor::ProcessSample(spSample);
		}

		_samplesInFragment++;
	}

	return hr;
}

HRESULT AudioSampleProcessor::InitializeFragments()
{
	HRESULT hr = ISampleProcessor::InitializeFragments(_fragmentIndex);

	if (SUCCEEDED(hr))
	{
		_currentFragmentDuration = 0;
		_currentFragmentOffset = _nextSampleTime;
		_fragmentEndTime = _currentFragmentOffset + _maxFragmentDuration;
		_isInitialized = true;
	}

	return hr;
}

HRESULT AudioSampleProcessor::TerminateFragments()
{
	HRESULT hr = ISampleProcessor::TerminateFragments(&_currentFragmentOffset, &_nextSampleTime);

	_fragmentIndex++;
	_samplesInFragment = 0;
	_isInitialized = false;

	return hr;
}